import numpy as np
import pandas as pd
import statsmodels.api as sm
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.stattools import adfuller
from statsmodels.stats.diagnostic import acorr_ljungbox
from statsmodels.tsa.arima.model import ARIMA
from data_loader import Dataset
import matplotlib.pyplot as plt
from metrics import metric, MSE
import os
from utils import savePlotResults, root_list, data_list
np.random.seed(1024)


class Arima:
    def __init__(self, root_path='./dataset/Brent', data_path='Brenta.csv'):
        self.root_path = root_path
        self.data_path = data_path
        self.dataset = Dataset(root_path, data_path)
        self.dtframe = self.__getFrame__()
        self.diff_data = {"1": None, "2": None}

    def __getFrame__(self):
        index = self.dataset.year
        frame = pd.DataFrame(self.dataset.data, index=index, columns=['price'])
        return frame

    def plot_frame(self):
        self.dtframe.plot(figsize=(20, 5))
        plt.show()

    def plot_cf(self, data):
        fig = plt.figure(figsize=(15, 8))
        ax1 = fig.add_subplot(211)
        plot_acf(data, ax=ax1)

        ax2 = fig.add_subplot(212)
        plot_pacf(data, ax=ax2)
        plt.show()

    def get_diff(self, diff: int = 1, plot=False):
        diff_data = self.dtframe.diff(diff).dropna()
        if plot:
            plt.title(f'diff={diff}')
            plt.plot(diff_data.index, diff_data.values)
            plt.show()
        self.diff_data[str(diff)] = diff_data

    def stable_judge(self, data):
        # result: (adf, pvalue, usedlag, nobs, critical_values, icbest)
        result = adfuller(data)
        index = ['adf', 'pvalue', 'usedlag', 'nobs', 'critical_values', 'icbest']
        for i, key in enumerate(index):
            print(f'{key}: {result[i]}')
        return result

    def white_noise_judge(self, data, lags=1):
        result = acorr_ljungbox(data, lags=lags)
        print(result)
        return result

    def getModel(self, p, d, q):
        train = np.concatenate([self.dataset.train_data.flatten(), self.dataset.val_data.flatten()])
        self.model = ARIMA(train, order=(p, d, q))
        res = self.model.fit()
        return res.summary()

    def predict(self, time_steps=1):
        assert hasattr(self, 'model')
        preds = self.model.fit().forecast(time_steps)
        return preds



if __name__ == '__main__':
    option = -2
    root_path = os.path.join('./dataset', root_list[option])
    data_path = data_list[option] + '.csv'
    """Model"""
    model = Arima(root_path, data_path)
    # model.getModel(3, 0, 6)   # Brent
    # model.getModel(1, 0, 2)   # RCLC
    # model.getModel(1, 0, 6)   # WTI
    # model.getModel(1, 0, 2)     # NYMEX
    model.getModel(1, 1, 2)  # OPEC
    """Diff"""
    model.get_diff(diff=1, plot=True)

    """ACF and PACF"""
    # model.plot_cf(model.diff_data['1'])

    """Stable"""
    # results1 = model.stable_judge(model.diff_data['1'])

    """white"""
    # results2 = model.white_noise_judge(model.diff_data['1'])

    dtframe = savePlotResults(-2, True, True, True, True, model, pred_steps=10, units="美元/百万英热")
